package com.wdemo.embedding;

import com.wdemo.enums.ChatTypeEnum;
import com.wdemo.enums.EmbeddingTypeEnum;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.stereotype.Component;
import redis.clients.jedis.JedisPooled;

import java.util.HashMap;
import java.util.Map;

/**
 * @ClassName ChatMemoryManager
 * @Description 向量库管理器
 * @Author WDEMO
 * @Date 2025/8/31 0:38
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class EmbeddingManager {

    private final JedisPooled jedisPooled;

    private final EmbeddingModel embeddingModel;


    private static final Map<EmbeddingTypeEnum, VectorStore> vectorStoreMap = new HashMap<>();

//    public ChatMemoryManager() {
//        log.info("初始化聊天内存管理器");
//        // 初始化不同类型的聊天内存
//        for (ChatTypeEnum type : ChatTypeEnum.values()) {
//            chatMemoryMap.put(type, new RedisChatMemory(type.getType(),stringRedisTemplate));
//        }
//        log.info("初始化聊天内存管理器完成");
//    }

    @PostConstruct
    public void init() {
        log.info("初始化向量库管理器");
        // 初始化不同类型的聊天内存
        for (EmbeddingTypeEnum type : EmbeddingTypeEnum.values()) {
            // 传递 StringRedisTemplate 给 RedisChatMemory
            vectorStoreMap.put(
                    type,
                    RedisVectorStore.builder(jedisPooled, embeddingModel)
                            .prefix(type.getType())
                            .indexName(type.getIndex())
                            .initializeSchema(true)
                            .build());
        }
        log.info("初始化向量库管理器完成");
    }

    public VectorStore getVectorStore(EmbeddingTypeEnum chatType) {
        return vectorStoreMap.get(chatType);
    }

    public VectorStore getVectorStore(String chatType) {
        return vectorStoreMap.get(ChatTypeEnum.fromString(chatType));
    }
}
